import numpy as np
import cv2 as cv
from matplotlib import pyplot as plt


def draw_circles(ex, img):
    # 用霍夫变换画圆
    circles = cv.HoughCircles(ex, cv.HOUGH_GRADIENT, 1, 10, param1=50, param2=9, minRadius=4, maxRadius=10)
    circles = np.uint16(np.around(circles))
    circles = circles[0, :]

    # 找出在同一行的点(尽量避免图片倾斜角度的干扰),并根据坐标轴对点进行排序
    c = 0
    l = circles.shape[0]
    grouped_circles = []
    while c < l:
        d = c + 1
        while d < l:
            if np.abs(int(circles[c, 1]) - int(circles[d, 1])) <= 17:
                grouped_circles.append([circles[c], circles[d]])
            d = d + 1
        c = c + 1
    grouped_circles = np.array(grouped_circles)
    sorted_grouped_circles = []
    for a in grouped_circles:
        a = a[np.lexsort((a[:, 1], a[:, 0]))]
        sorted_grouped_circles.append(a)
    sorted_grouped_circles = np.array(sorted(sorted_grouped_circles, key=lambda x: x[0][1]))
    sorted_grouped_circles = sorted_grouped_circles.reshape(-1, sorted_grouped_circles.shape[-1])

    j = 0
    for i in sorted_grouped_circles:
        # 查看圆心周围像素点值,一旦有255便判定为空心
        a = np.int8(-2)
        while 1:
            if ex[i[1] + a, i[0] + a] == 255:
                cv.circle(img, (i[0], i[1]), i[2], (0, 255, 0), 2)
                cv.putText(img, '0,%d' % j, (i[0], i[1] + 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                j = j + 1
                break
            a = a + 1
            if a == 3:
                cv.circle(img, (i[0], i[1]), i[2], (0, 0, 255), 2)
                cv.putText(img, '1,%d' % j, (i[0], i[1] + 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                j = j + 1
                break


if __name__ == '__main__':
    # 读图/转为灰度图
    img1 = cv.imread('ex1.jpg')
    ex1 = cv.cvtColor(img1, cv.COLOR_BGR2GRAY)
    img2 = cv.imread('ex2.jpg')
    ex2 = cv.cvtColor(img2, cv.COLOR_BGR2GRAY)
    # 进行高斯模糊 核值小时会导致实心变空心
    ex1 = cv.GaussianBlur(ex1, (9, 9), 0)
    ex2 = cv.GaussianBlur(ex2, (9, 9), 0)
    # 自适应阈值(均值) 使用高斯算法时也会导致实心变空心
    ex1 = cv.adaptiveThreshold(ex1, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 10)
    ex2 = cv.adaptiveThreshold(ex2, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY, 11, 10)
    # 初始化SIFT
    sift = cv.SIFT.create()
    kp1, des1 = sift.detectAndCompute(ex1, None)
    kp2, des2 = sift.detectAndCompute(ex2, None)
    # 蛮力匹配
    bf = cv.BFMatcher()
    matches = bf.knnMatch(des1, des2, k=2)
    # 用ratio test筛选得到的匹配
    good_matches = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good_matches.append(m)

    # 获取旋转角度
    rotation_angle = 0
    assert len(good_matches) > 4, "匹配点数量不足，请调整匹配参数"  # 确保有足够的匹配点
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    H, mask = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)  # 得到单应性矩阵H
    rotation_matrix = H[:, :2]  # 获取旋转部分矩阵（前两列）
    rotation_angle = np.degrees(np.arctan2(rotation_matrix[1, 0], rotation_matrix[0, 0]))

    # 旋转还原图2
    rows, cols = ex2.shape
    M = cv.getRotationMatrix2D(((cols - 1) / 2.0, (rows - 1) / 2.0), rotation_angle, 1)
    rotated_img2 = cv.warpAffine(img2, M, (cols, rows))
    rotated_ex2 = cv.warpAffine(ex2, M, (cols, rows))

    # 调用创建的draw_circles函数画圆
    draw_circles(ex1, img1)
    draw_circles(rotated_ex2, rotated_img2)

    # 改成rgb格式
    b, g, r = cv.split(img1)
    img1 = cv.merge([r, g, b])
    b, g, r = cv.split(rotated_img2)
    rotated_img2 = cv.merge([r, g, b])

    plt.subplot(1, 2, 1), plt.imshow(img1), plt.xticks([]), plt.yticks([])
    plt.subplot(1, 2, 2), plt.imshow(rotated_img2), plt.xticks([]), plt.yticks([])

    plt.show()
